package com.ark.common.mybatis.interceptor;

import com.ark.common.mybatis.permission.DataPermissionRule;
import com.ark.common.mybatis.permission.expression.AndExpressionX;
import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.baomidou.mybatisplus.extension.parser.JsqlParserSupport;
import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor;
import lombok.RequiredArgsConstructor;
import net.sf.jsqlparser.expression.Alias;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.ExpressionVisitorAdapter;
import net.sf.jsqlparser.expression.NullValue;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SelectItem;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.util.CollectionUtils;

import java.sql.SQLException;
import java.util.List;

@RequiredArgsConstructor
public class DataPermissionInterceptor extends JsqlParserSupport implements InnerInterceptor {

    private final DataPermissionRule rule;

    @Override
    public void beforeQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) throws SQLException {
        //调用父类JsqlParserSupport的processSelect方法，将SQL语句解析出来
        PluginUtils.MPBoundSql mpBs = PluginUtils.mpBoundSql(boundSql);
        //获取新的SQL 并设置到BoundSql对象中
        mpBs.sql(this.parserSingle(mpBs.sql(), ms.getId()));
    }

    /**
     * 处理查询语句
     */
    @Override
    protected void processSelect(Select select, int index, String sql, Object obj) {
        PlainSelect plainSelect = select.getPlainSelect();
        //向上转型，转为table获取表名
        Table table = (Table) plainSelect.getFromItem();
        if (CollectionUtils.isEmpty(rule.getTableNames()) || !rule.getTableNames().contains(table.getName())) {
            return;
        }
        //获取所有select查询的字段
        List<SelectItem<?>> selectItems = plainSelect.getSelectItems();
        Alias alias = table.getAlias();
        if (alias == null) {
            //设置别名,如果没有别名，则设置别名为表名的首字母
            //TODO : 如果是 xxx_xxx的中间表，关联表，则别名为别名应该取每段的首字母
            table.setAlias(new Alias(table.getName().substring(0, 1)));
            //遍历selectItems，设置字段为表别名+字段名
            selectItems.forEach(selectItem -> {
                //获取具体的表达式expression，其实就是查询的字段Column
                Expression expression = selectItem.getExpression();
                //修改查询字段表达式为表别名+字段
                if (expression instanceof Column) {
                    Column column = (Column) expression;
                    column.setTable(table);
                }
            });
        }
        Expression where = plainSelect.getWhere();
        //修改原where条件为表别名+字段，避免多表联查下的错误
        Expression expression = rule.getExpression(table);

        if (expression instanceof NullValue) return;
        if (where != null) {
            where.accept(new ExpressionVisitorAdapter() {
                @Override
                public void visit(Column column) {
                    column.setTable(table);
                }
            });
            AndExpression andExpression = new AndExpressionX();
            andExpression.setLeftExpression(where);
            andExpression.setRightExpression(expression);
            plainSelect.setWhere(andExpression);
            return;
        }
        plainSelect.setWhere(expression);
    }
}
